package com.google.code.maven2.plugins.hotswap;

import com.sun.jdi.Bootstrap;
import com.sun.jdi.ReferenceType;
import com.sun.jdi.VirtualMachine;
import com.sun.jdi.VirtualMachineManager;
import com.sun.jdi.connect.AttachingConnector;
import com.sun.jdi.connect.Connector;
import com.sun.jdi.connect.Transport;
import org.apache.maven.plugin.logging.Log;

import java.io.DataInputStream;
import java.io.IOException;
import java.net.URL;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


/**
 * This class provides the workings necessary to connect to a running JVM and to replace classes.
 *
 * This class was taken from http://code.google.com/p/hotswap/source/browse/trunk/java/dak/debug/HotSwapHelper.java?spec=svn3&r=3 then modified by Mike Nereson
 *
 * @author David A. Kavanagh <a href="mailto:dak@dotech.com">dak@dotech.com</a>
 */
public class HotSwapHelper
{
    /**
     * The virtual machine that is connected to from the given host and port.
     */
    private VirtualMachine virtualMachine;

    /**
     * Reference to the Maven log
     */
    final protected Log log;

    private static final String DT_SHMEM = "dt_shmem";
    private static final String DT_SOCKET = "dt_socket";
    private static final String PORT = "port";
    private static final String HOST = "hostname";
    private static final String NAME = "name";

    public HotSwapHelper(final Log log)
    {
        this.log = log;
    }

    public void connect(String name)
            throws Exception
    {
        VirtualMachineManager virtualMachineManager = Bootstrap.virtualMachineManager();
        List<AttachingConnector> connectors = virtualMachineManager.attachingConnectors();

        AttachingConnector shmemAttachingConnector = null;

        for (AttachingConnector connector : connectors)
        {
            Transport transport = connector.transport();

            if (DT_SHMEM.equals(transport.name()))
            {
                shmemAttachingConnector = (AttachingConnector) connector;
                break;
            }
        }


        if (null == shmemAttachingConnector)
        {
            throw new RuntimeException("Shared memory attaching connector not found");
        }

        Map<String, ? extends Connector.Argument> arguments = shmemAttachingConnector.defaultArguments();
        Connector.Argument memArg = (Connector.Argument) arguments.get(NAME);

        memArg.setValue(name);

        this.virtualMachine = shmemAttachingConnector.attach(arguments);

        if (!this.virtualMachine.canRedefineClasses())
        {
            throw new RuntimeException("JVM doesn't support class replacement");
        }
    }

    // either host, port will be set, or name
    public void connect(final String host, final String port)
            throws Exception
    {
        // connect to JVM
        VirtualMachineManager virtualMachineManager = Bootstrap.virtualMachineManager();
        List<AttachingConnector> connectors = virtualMachineManager.attachingConnectors();

        AttachingConnector socketAttachingConnector = null;

        for (AttachingConnector connector : connectors)
        {
            Transport transport = connector.transport();

            if (DT_SOCKET.equals(transport.name()))
            {
                socketAttachingConnector = (AttachingConnector) connector;
                break;
            }
        }

        if (null == socketAttachingConnector)
        {
            throw new RuntimeException("Socket attaching connector not found");
        }

        Map<String, ? extends Connector.Argument> arguments = socketAttachingConnector.defaultArguments();
        Connector.Argument hostArg = (Connector.Argument) arguments.get(HOST);
        Connector.Argument portArg = (Connector.Argument) arguments.get(PORT);

        hostArg.setValue(host);
        portArg.setValue(String.valueOf(port));

        this.virtualMachine = socketAttachingConnector.attach(arguments);

        if (!this.virtualMachine.canRedefineClasses())
        {
            throw new RuntimeException("JVM doesn't support class replacement");
        }
    }

    public void disconnect()
            throws Exception
    {
    }

    /**
     * TODO: this method swaps a single class at at time, however, VirtualMachine#redefineClasses takes a map, so all the swaps could be done at once.
     *
     * Is the performance hit of individual swaps wroth switching? note that the current one-at-a-time method allows for reporting which class failed
     *
     * @param classFile
     * @param className
     * @throws Exception
     */
    public void replace(final URL classFile, final String className)
            throws Exception
    {
        // load class(es)
        final byte[] classBytes = loadClassFile(classFile);

        // redefine in JVM
        final List classes = virtualMachine.classesByName(className);

        // if the class isn't loaded on the VM, can't do the replace.
        if (classes == null || classes.isEmpty())
        {
            return;
        }

        final HashMap map = new HashMap();

        for (int i = 0, length = classes.size(); i < length; i++)
        {
            final ReferenceType referenceType = (ReferenceType) classes.get(i);
            map.put(referenceType, classBytes);

            try
            {
                virtualMachine.redefineClasses(map);
            }
            catch (Exception e)
            {
                final String classname = referenceType.signature().substring(1, referenceType.signature().length() - 1);
                final String cause = e.getMessage();

                final String message = "Failed to swap " + classname + " because " + cause;

                log.warn(message);

                throw new RuntimeException(message);
            }

            log.debug("\t\t\t\tHotswapped " + referenceType.signature().substring(1));
        }
    }

    private byte[] loadClassFile(final URL classFileURL)
            throws IOException
    {
        final DataInputStream inputStream = new DataInputStream(classFileURL.openStream());

        final int length = classFileURL.openConnection().getContentLength();
        final byte[] bytes = new byte[length];

        inputStream.readFully(bytes);
        inputStream.close();

        return bytes;
    }
}
